package com.lsa.helpers.heap;

import java.util.Arrays;
import java.util.Comparator;

import com.lsa.helpers.Utils;

public class ArrayHeap<T> {
	
	private static final int INITIAL_SIZE = 10;
	private Comparator<T> comparator;
	private Object[] data;
	private int size = 0;
	
	public ArrayHeap(Comparator<T> comparator) {
		this.comparator = comparator;
	}
	
	public void insert(T element){
		if (size == 0){
			data = new Object[INITIAL_SIZE];
		}
		else if (size == data.length){
			Object[] newData = new Object[data.length * 2];
			System.arraycopy(data, 0, newData, 0, data.length);
			data = newData;
		}
		data[size] = element;
		
		for(int i = size + 1; i > 1; ){
			int parentIndex = getParentIndex(i);
			if (comparator.compare(get(i), get(parentIndex)) <= 0){ // at correct place
				break;
			}
			else{
				Utils.swap(data, i - 1, parentIndex - 1);
			}
			i = parentIndex;
		}
		
		++size;
	}
	
	public T pop(){
		T returnValue = get(1);
		if (getSize() > 1){
			set(1, get(size));
		}
		--size;
		heapify(1, size);
		
		return returnValue;
	}
	
	public int getSize() {
		return size;
	}
	
	private void heapify(int rootIndex, int length){
		if (size == 0) return;		
		T rootElement = get(rootIndex);
		
		T leftSubtreeRootElement = getLeftSubtree(rootIndex, length);
		T rightSubtreeRootElement = getRightSubtree(rootIndex, length);
		
		if (leftSubtreeRootElement == null && rightSubtreeRootElement == null){
			return;
		}
		
		if ( rightSubtreeRootElement == null){
			if (comparator.compare(rootElement ,leftSubtreeRootElement) < 0){
				swapWithSubtree(rootIndex, rootIndex * 2);
			}
			heapify(rootIndex * 2, length);
			return;
		}
		
		int maxSubtreeIndex = maxSubtreeIndex(rootIndex);
		if (comparator.compare(get(maxSubtreeIndex), rootElement) > 0){
			swapWithSubtree(rootIndex, maxSubtreeIndex);
			heapify(maxSubtreeIndex, length);
		}
	}
	
	private void swapWithSubtree(int root, int subtreeRootIndex){
		T rootElement = get(root);
		set(root, get(subtreeRootIndex));
		set(subtreeRootIndex, rootElement);
	}
	
	private int maxSubtreeIndex(int root){
		int leftSubtreeIndex = root * 2;
		int rightSubtreeIndex = root * 2 + 1;
		
		if (comparator.compare(get(leftSubtreeIndex), get(rightSubtreeIndex)) < 0) return rightSubtreeIndex;
		
		return leftSubtreeIndex;
	}
	
	private T getLeftSubtree(int root, int length){
		int subtreeRootIndex = root * 2;
		if (subtreeRootIndex > length ){
			return null;
		}
		
		return get(subtreeRootIndex);
	}
	
	private T getRightSubtree(int root, int length){
		int subtreeRootIndex = root * 2 + 1;
		if (subtreeRootIndex > length ){
			return null;
		}
		
		return get(subtreeRootIndex);
	}

	private int getParentIndex(int i){
		return i / 2;
	}
	
	private T get(int i){
		return (T)data[i - 1];
	}
	
	private void set(int i, T elem){
		data[i - 1] = elem;
	}
	
	@SuppressWarnings("unchecked")
	public <T> T[] toArray(T[] a) {
		if (a.length < size)
			return (T[]) Arrays.copyOf(data, size, a.getClass());
		System.arraycopy(data, 0, a, 0, size);
		if (a.length > size)
			a[size] = null;
		return a;
	}
}
